struct ZUCLFSR {
    lfsr: [u32; 16],
}

impl ZUCLFSR {
    fn new(lfsr: [u32; 16]) -> Self {
        Self {
            lfsr,
        }
    }
    fn mulpow2mod<const K: u32>(x: u32) -> u32 {
        ((x << K) | (x >> (31 - K))) & 0x7FFFFFFF
    }

    fn addmod(x: u32, y: u32) -> u32 {
        let ret = x+y;
        let ret = (ret & 0x7FFFFFFF) + (ret >> 31);
        ret
    }
    
    fn with_initialisation_mode(&mut self, u: u32) {
        let mut tmp = self.lfsr[0];
        tmp = Self::addmod(tmp, Self::mulpow2mod::<8>(self.lfsr[0]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<20>(self.lfsr[4]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<21>(self.lfsr[10]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<17>(self.lfsr[13]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<15>(self.lfsr[15]));
        tmp = Self::addmod(tmp, u);
        self.lfsr.copy_within(1..16, 0);
        self.lfsr[15] = tmp;
    }

    fn with_work_mode(&mut self) {
        let mut tmp = self.lfsr[0];
        tmp = Self::addmod(tmp, Self::mulpow2mod::<8>(self.lfsr[0]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<20>(self.lfsr[4]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<21>(self.lfsr[10]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<17>(self.lfsr[13]));
        tmp = Self::addmod(tmp, Self::mulpow2mod::<15>(self.lfsr[15]));
        self.lfsr.copy_within(1..16, 0);
        self.lfsr[15] = tmp;
    }

    fn bit_reorganization(&self) -> (u32, u32, u32, u32) {
        (
            ((self.lfsr[15] & 0x7FFF8000) << 1) | (self.lfsr[14] & 0xFFFF),
            ((self.lfsr[11] & 0xFFFF) << 16) | (self.lfsr[9] >> 15),
            ((self.lfsr[7] & 0xFFFF) << 16) | (self.lfsr[5] >> 15),
            ((self.lfsr[2] & 0xFFFF) << 16) | (self.lfsr[0] >> 15),
        )
    }
}


struct ZUCNonlinearFunction {
    r1: u32,
    r2: u32,
}

impl ZUCNonlinearFunction {
    fn new() -> Self {
        Self {
            r1: 0,
            r2: 0,
        }
    }
    fn s(x: u32) -> u32 {
        const S0: [u8; 256] = [
            0x3E,0x72,0x5B,0x47,0xCA,0xE0,0x00,0x33,0x04,0xD1,0x54,0x98,0x09,0xB9,0x6D,0xCB,
            0x7B,0x1B,0xF9,0x32,0xAF,0x9D,0x6A,0xA5,0xB8,0x2D,0xFC,0x1D,0x08,0x53,0x03,0x90,
            0x4D,0x4E,0x84,0x99,0xE4,0xCE,0xD9,0x91,0xDD,0xB6,0x85,0x48,0x8B,0x29,0x6E,0xAC,
            0xCD,0xC1,0xF8,0x1E,0x73,0x43,0x69,0xC6,0xB5,0xBD,0xFD,0x39,0x63,0x20,0xD4,0x38,
            0x76,0x7D,0xB2,0xA7,0xCF,0xED,0x57,0xC5,0xF3,0x2C,0xBB,0x14,0x21,0x06,0x55,0x9B,
            0xE3,0xEF,0x5E,0x31,0x4F,0x7F,0x5A,0xA4,0x0D,0x82,0x51,0x49,0x5F,0xBA,0x58,0x1C,
            0x4A,0x16,0xD5,0x17,0xA8,0x92,0x24,0x1F,0x8C,0xFF,0xD8,0xAE,0x2E,0x01,0xD3,0xAD,
            0x3B,0x4B,0xDA,0x46,0xEB,0xC9,0xDE,0x9A,0x8F,0x87,0xD7,0x3A,0x80,0x6F,0x2F,0xC8,
            0xB1,0xB4,0x37,0xF7,0x0A,0x22,0x13,0x28,0x7C,0xCC,0x3C,0x89,0xC7,0xC3,0x96,0x56,
            0x07,0xBF,0x7E,0xF0,0x0B,0x2B,0x97,0x52,0x35,0x41,0x79,0x61,0xA6,0x4C,0x10,0xFE,
            0xBC,0x26,0x95,0x88,0x8A,0xB0,0xA3,0xFB,0xC0,0x18,0x94,0xF2,0xE1,0xE5,0xE9,0x5D,
            0xD0,0xDC,0x11,0x66,0x64,0x5C,0xEC,0x59,0x42,0x75,0x12,0xF5,0x74,0x9C,0xAA,0x23,
            0x0E,0x86,0xAB,0xBE,0x2A,0x02,0xE7,0x67,0xE6,0x44,0xA2,0x6C,0xC2,0x93,0x9F,0xF1,
            0xF6,0xFA,0x36,0xD2,0x50,0x68,0x9E,0x62,0x71,0x15,0x3D,0xD6,0x40,0xC4,0xE2,0x0F,
            0x8E,0x83,0x77,0x6B,0x25,0x05,0x3F,0x0C,0x30,0xEA,0x70,0xB7,0xA1,0xE8,0xA9,0x65,
            0x8D,0x27,0x1A,0xDB,0x81,0xB3,0xA0,0xF4,0x45,0x7A,0x19,0xDF,0xEE,0x78,0x34,0x60,
        ];
        const S1: [u8; 256] = [
            0x55,0xC2,0x63,0x71,0x3B,0xC8,0x47,0x86,0x9F,0x3C,0xDA,0x5B,0x29,0xAA,0xFD,0x77,
            0x8C,0xC5,0x94,0x0C,0xA6,0x1A,0x13,0x00,0xE3,0xA8,0x16,0x72,0x40,0xF9,0xF8,0x42,
            0x44,0x26,0x68,0x96,0x81,0xD9,0x45,0x3E,0x10,0x76,0xC6,0xA7,0x8B,0x39,0x43,0xE1,
            0x3A,0xB5,0x56,0x2A,0xC0,0x6D,0xB3,0x05,0x22,0x66,0xBF,0xDC,0x0B,0xFA,0x62,0x48,
            0xDD,0x20,0x11,0x06,0x36,0xC9,0xC1,0xCF,0xF6,0x27,0x52,0xBB,0x69,0xF5,0xD4,0x87,
            0x7F,0x84,0x4C,0xD2,0x9C,0x57,0xA4,0xBC,0x4F,0x9A,0xDF,0xFE,0xD6,0x8D,0x7A,0xEB,
            0x2B,0x53,0xD8,0x5C,0xA1,0x14,0x17,0xFB,0x23,0xD5,0x7D,0x30,0x67,0x73,0x08,0x09,
            0xEE,0xB7,0x70,0x3F,0x61,0xB2,0x19,0x8E,0x4E,0xE5,0x4B,0x93,0x8F,0x5D,0xDB,0xA9,
            0xAD,0xF1,0xAE,0x2E,0xCB,0x0D,0xFC,0xF4,0x2D,0x46,0x6E,0x1D,0x97,0xE8,0xD1,0xE9,
            0x4D,0x37,0xA5,0x75,0x5E,0x83,0x9E,0xAB,0x82,0x9D,0xB9,0x1C,0xE0,0xCD,0x49,0x89,
            0x01,0xB6,0xBD,0x58,0x24,0xA2,0x5F,0x38,0x78,0x99,0x15,0x90,0x50,0xB8,0x95,0xE4,
            0xD0,0x91,0xC7,0xCE,0xED,0x0F,0xB4,0x6F,0xA0,0xCC,0xF0,0x02,0x4A,0x79,0xC3,0xDE,
            0xA3,0xEF,0xEA,0x51,0xE6,0x6B,0x18,0xEC,0x1B,0x2C,0x80,0xF7,0x74,0xE7,0xFF,0x21,
            0x5A,0x6A,0x54,0x1E,0x41,0x31,0x92,0x35,0xC4,0x33,0x07,0x0A,0xBA,0x7E,0x0E,0x34,
            0x88,0xB1,0x98,0x7C,0xF3,0x3D,0x60,0x6C,0x7B,0xCA,0xD3,0x1F,0x32,0x65,0x04,0x28,
            0x64,0xBE,0x85,0x9B,0x2F,0x59,0x8A,0xD7,0xB0,0x25,0xAC,0xAF,0x12,0x03,0xE2,0xF2,
        ];
        ((S0[(x >> 24) as usize] as u32) << 24)
        | ((S1[((x >> 16) & 0xFF) as usize] as u32) << 16)
        | ((S0[((x >> 8) & 0xFF) as usize] as u32) << 8)
        | (S1[(x & 0xFF) as usize] as u32)
    }

    fn rot32<const K: u32>(x: u32) -> u32 {
        (x << K) | (x >> (32 - K))
    }

    fn l1(x: u32) -> u32 {
        x
        ^ Self::rot32::<2>(x)
        ^ Self::rot32::<10>(x)
        ^ Self::rot32::<18>(x)
        ^ Self::rot32::<24>(x)
    }

    fn l2(x: u32) -> u32 {
        x
        ^ Self::rot32::<8>(x)
        ^ Self::rot32::<14>(x)
        ^ Self::rot32::<22>(x)
        ^ Self::rot32::<30>(x)
    }

    fn f(&mut self, x: (u32, u32, u32)) -> u32 {
        let w = (x.0 ^ self.r1).overflowing_add(self.r2).0;
        let w1 = self.r1.overflowing_add(x.1).0;
        let w2 = self.r2 ^ x.2;
        self.r1 = Self::s(Self::l1((w1 << 16) | (w2 >> 16)));
        self.r2 = Self::s(Self::l2((w2 << 16) | (w1 >> 16)));
        w
    }
}

struct ZUC {
    lfsr: ZUCLFSR,
    f: ZUCNonlinearFunction,
}

impl ZUC {
    fn make_u31(a: u8, b: u32, c: u8) -> u32 {
        ((a as u32) << 23) | (b << 8) | (c as u32)
    }

    pub fn new(k: &[u8], iv: &[u8]) -> Self {
        const EK_D: [u32; 16] = [
            0x44D7, 0x26BC, 0x626B, 0x135E, 0x5789, 0x35E2, 0x7135, 0x09AF,
            0x4D78, 0x2F13, 0x6BC4, 0x1AF1, 0x5E26, 0x3C4D, 0x789A, 0x47AC,
        ];
        let lfsr = (0..16).map(|i|
            Self::make_u31(k[i], EK_D[i], iv[i]))
            .collect::<Vec<u32>>().try_into().unwrap();
        let mut lfsr = ZUCLFSR::new(lfsr);
        let mut f = ZUCNonlinearFunction::new();
        for _ in 0..32 {
            let (x0, x1, x2, _) = lfsr.bit_reorganization();
            let w = f.f((x0, x1, x2));
            lfsr.with_initialisation_mode(w >> 1);
        }
        let (x0, x1, x2, _) = lfsr.bit_reorganization();
        let _ = f.f((x0, x1, x2));
        lfsr.with_work_mode();
        Self {
            lfsr,
            f,
        }
    }
    pub fn generate_keystream(&mut self, len: usize) -> Vec<u32> {
        (0..len).map(|_| {
            let (x0, x1, x2, x3) = self.lfsr.bit_reorganization();
            let ret = self.f.f((x0, x1, x2)) ^ x3;
            self.lfsr.with_work_mode();
            ret
        }).collect()
    }


}

fn pkcs2padding(s: &[u8]) -> Vec<u8> {
    const BLOCK_SIZE: usize = 4;
    let padding_len = BLOCK_SIZE - (s.len() % BLOCK_SIZE); 
    let mut ret = Vec::with_capacity(s.len() + padding_len);
    ret.extend_from_slice(s);
    ret.extend(vec![padding_len as u8; padding_len]);
    ret
}

pub fn encryption(input: String) -> String {
    let mut zuc = ZUC::new(&[0u8; 16], &[1u8; 16]);
    let input = pkcs2padding(input.as_bytes());
    let len = input.len();
    let stream: Vec<u8> = zuc.generate_keystream(len / 4).iter().flat_map(|&x|{
        vec![
            
            
            
            ((x >> 24) & 0xFF) as u8,
            ((x >> 16) & 0xFF) as u8,
            ((x >> 8) & 0xFF) as u8,
            (x & 0xFF) as u8,
        ]
    }).collect();
    let result: Vec<u8> = (0..len).map(|i| {
        input[i] ^ stream[i]
    }).collect();    
    base64::Engine::encode(&base64::prelude::BASE64_STANDARD, result)
}
